#pragma once
#include <cinttypes>
#include <type_traits>

#include "container.h"
#include "math_util.h"

template <typename Element_, int NumChannelsPacked, int AccessSize,
          int NumThreadsCTA = 128, int MaxNumScales = 32>
struct ParamsBase {
  using Element = Element_;
  void *data_output;  // [N,Q,Head,C]
  utils::Array<const void *, MaxNumScales>
      data_value;                     // [N,Head*C/x,H,W,x] or [N,H,W,Head*C]
  const void *data_sampling_offsets;  // [N,Q,Head,S,P,2]
  const void *data_attn_weight;       // [N,Q,Head,S,P]
  utils::Array<int2, MaxNumScales> spatial_shapes;   // [S,2]
  utils::Array<float2, MaxNumScales> qscale_offset;  // [S,2]
  utils::Array<float, MaxNumScales> qscale_i2o;      // [S]
  FastDivmod div_mod_channels_per_head;
  int BatchSize;
  int NumQueries;
  int NumHeads;
  int NumChannelsPerHead;
  int NumScales;
  int NumPointsPerScale;
  int NumChannels;
  int NumPoints;

  static_assert((AccessSize * sizeof(Element)) % 4 == 0,
                "AccessSize must be multiples of 4 bytes");
  static_assert(NumChannelsPacked % AccessSize == 0,
                "NumChannelsPacked must be in multiples of AccessSize");

  static constexpr int kNumChannelsPacked = NumChannelsPacked;
  static constexpr int kAccessSize = AccessSize;
  static constexpr int kAccessSize4B = (AccessSize * sizeof(Element)) / 4;

  static constexpr int kNumThreadsCTA = NumThreadsCTA;
  static constexpr int kNumThreadsContiguousCTA =
      kNumChannelsPacked / kAccessSize;
  static constexpr int kNumThreadsStridedCTA =
      kNumThreadsCTA / kNumThreadsContiguousCTA;

  static_assert(kNumThreadsContiguousCTA > 0,
                "kNumThreadsContiguousCTA must be greater than 0");
  static_assert(kNumThreadsStridedCTA > 0,
                "kNumThreadsStridedCTA must be greater than 0");

  HOST_DEVICE_INLINE ParamsBase(
      void *data_output, const void **data_multiscale_value,
      const void *data_sampling_offsets, const void *data_attn_weight,
      const int2 *multiscale_spatial_shapes, const int BatchSize,
      const int NumQueries, const int NumHeads, const int NumScales,
      const int NumPointsPerScale, const int NumChannelsPerHead,
      const float *qscale_value, const float qscale_offset,
      const float qscale_weight, const float qscale_out)
      : data_output(data_output),
        data_sampling_offsets(data_sampling_offsets),
        data_attn_weight(data_attn_weight),
        div_mod_channels_per_head(NumChannelsPerHead),
        BatchSize(BatchSize),
        NumHeads(NumHeads),
        NumChannelsPerHead(NumChannelsPerHead),
        NumQueries(NumQueries),
        NumScales(NumScales),
        NumPointsPerScale(NumPointsPerScale),
        NumPoints(NumScales * NumPointsPerScale),
        NumChannels(NumHeads * NumChannelsPerHead) {
    for (int s = 0; s < NumScales; s++) {
      this->data_value[s] = data_multiscale_value[s];
      this->spatial_shapes[s].x = multiscale_spatial_shapes[s].y;
      this->spatial_shapes[s].y = multiscale_spatial_shapes[s].x;
      this->qscale_offset[s].x = qscale_offset * this->spatial_shapes[s].x;
      this->qscale_offset[s].y = qscale_offset * this->spatial_shapes[s].y;
      this->qscale_i2o[s] = qscale_weight * qscale_value[s] / qscale_out;
    }
  }

  /** |<--------- loop in CTA --------->|
   *      threadIdx.x        threadIdx.y
   *  heads*channels(%32) ---> queries ---> heads*channels(/32) ---> batch
   *                          blockIdx.x        blockIdx.y        blockIdx.z
   *                        |<---------- loop with multiple CTAs ---------->| */

  inline dim3 block_size() const {
    return dim3(kNumThreadsContiguousCTA, kNumThreadsStridedCTA, 1);
  }
  inline dim3 grid_size() const {
    return dim3(div_up(NumQueries, kNumThreadsStridedCTA),
                div_up(NumHeads * NumChannelsPerHead, kNumChannelsPacked),
                BatchSize);
  }
};

template <typename Element, Layout SourceLayout, int AccessSize,
          int MaxNumScales = 32>
struct Params;

template <int AccessSize, int MaxNumScales>
struct Params<int8_t, Layout::NCHW32, AccessSize, MaxNumScales>
    : ParamsBase<int8_t, 32, AccessSize, 128, MaxNumScales> {
  static constexpr Layout kSourceLayout = Layout::NCHW32;
  using Base = ParamsBase<int8_t, 32, AccessSize, 128, MaxNumScales>;

  utils::Array<int, MaxNumScales> stride_height_;  // [S]
  utils::Array<int, MaxNumScales> stride_packed_;  // [S]
  utils::Array<int, MaxNumScales> stride_batch_;   // [S]

  template <typename... Args>
  HOST_DEVICE_INLINE Params(Args... args) : Base(args...) {
    for (int s = 0; s < this->NumScales; s++) {
      stride_height_[s] = this->spatial_shapes[s].x * Base::kNumChannelsPacked;
      stride_packed_[s] = this->spatial_shapes[s].y * stride_height_[s];
      stride_batch_[s] = div_up(this->NumHeads * this->NumChannelsPerHead,
                                Base::kNumChannelsPacked) *
                         stride_packed_[s];
    }
  }
  HOST_DEVICE_INLINE int stride_width(int idx_scale) const {
    return Base::kNumChannelsPacked;
  }
  HOST_DEVICE_INLINE int stride_height(int idx_scale) const {
    return stride_height_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_packed(int idx_scale) const {
    return stride_packed_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_batch(int idx_scale) const {
    return stride_batch_[idx_scale];
  }
};

template <int AccessSize, int MaxNumScales>
struct Params<int8_t, Layout::NHWC, AccessSize, MaxNumScales>
    : ParamsBase<int8_t, 128, AccessSize, 256, MaxNumScales> {
  static constexpr Layout kSourceLayout = Layout::NHWC;
  using Base = ParamsBase<int8_t, 128, AccessSize, 256, MaxNumScales>;

  int stride_width_;                               // [S]
  utils::Array<int, MaxNumScales> stride_height_;  // [S]
  utils::Array<int, MaxNumScales> stride_batch_;   // [S]

  template <typename... Args>
  HOST_DEVICE_INLINE Params(Args... args) : Base(args...) {
    stride_width_ = this->NumHeads * this->NumChannelsPerHead;
    for (int s = 0; s < this->NumScales; s++) {
      stride_height_[s] = this->spatial_shapes[s].x * stride_width_;
      stride_batch_[s] = this->spatial_shapes[s].y * stride_height_[s];
    }
  }
  HOST_DEVICE_INLINE int stride_width(int idx_scale) const {
    return stride_width_;
  }
  HOST_DEVICE_INLINE int stride_height(int idx_scale) const {
    return stride_height_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_packed(int idx_scale) const {
    return Base::kNumChannelsPacked;
  }
  HOST_DEVICE_INLINE int stride_batch(int idx_scale) const {
    return stride_batch_[idx_scale];
  }
};

template <int AccessSize, int MaxNumScales>
struct Params<half, Layout::NCHW32, AccessSize, MaxNumScales>
    : ParamsBase<half, 32, AccessSize, 128, MaxNumScales> {
  static constexpr Layout kSourceLayout = Layout::NCHW32;
  using Base = ParamsBase<half, 32, AccessSize, 128, MaxNumScales>;

  utils::Array<int, MaxNumScales> stride_height_;  // [S]
  utils::Array<int, MaxNumScales> stride_packed_;  // [S]
  utils::Array<int, MaxNumScales> stride_batch_;   // [S]

  utils::Array<half2, MaxNumScales> spatial_shapes_half2;  // [S,2]

  template <typename... Args>
  HOST_DEVICE_INLINE Params(Args... args) : Base(args...) {
    for (int s = 0; s < this->NumScales; s++) {
      spatial_shapes_half2[s] =
          half2{this->spatial_shapes[s].x, this->spatial_shapes[s].y};
      stride_height_[s] = this->spatial_shapes[s].x * Base::kNumChannelsPacked;
      stride_packed_[s] = this->spatial_shapes[s].y * stride_height_[s];
      stride_batch_[s] = div_up(this->NumHeads * this->NumChannelsPerHead,
                                Base::kNumChannelsPacked) *
                         stride_packed_[s];
    }
  }
  HOST_DEVICE_INLINE int stride_width(int idx_scale) const {
    return Base::kNumChannelsPacked;
  }
  HOST_DEVICE_INLINE int stride_height(int idx_scale) const {
    return stride_height_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_packed(int idx_scale) const {
    return stride_packed_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_batch(int idx_scale) const {
    return stride_batch_[idx_scale];
  }
};

template <int AccessSize, int MaxNumScales>
struct Params<half, Layout::NHWC, AccessSize, MaxNumScales>
    : ParamsBase<half, 64, AccessSize, 256, MaxNumScales> {
  static constexpr Layout kSourceLayout = Layout::NHWC;
  using Base = ParamsBase<half, 64, AccessSize, 256, MaxNumScales>;

  int stride_width_;                                       // [S]
  utils::Array<int, MaxNumScales> stride_height_;          // [S]
  utils::Array<int, MaxNumScales> stride_batch_;           // [S]
  utils::Array<half2, MaxNumScales> spatial_shapes_half2;  // [S,2]

  template <typename... Args>
  HOST_DEVICE_INLINE Params(Args... args) : Base(args...) {
    stride_width_ = this->NumHeads * this->NumChannelsPerHead;
    for (int s = 0; s < this->NumScales; s++) {
      spatial_shapes_half2[s] =
          half2{this->spatial_shapes[s].x, this->spatial_shapes[s].y};
      stride_height_[s] = this->spatial_shapes[s].x * stride_width_;
      stride_batch_[s] = this->spatial_shapes[s].y * stride_height_[s];
    }
  }
  HOST_DEVICE_INLINE int stride_width(int idx_scale) const {
    return stride_width_;
  }
  HOST_DEVICE_INLINE int stride_height(int idx_scale) const {
    return stride_height_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_packed(int idx_scale) const {
    return Base::kNumChannelsPacked;
  }
  HOST_DEVICE_INLINE int stride_batch(int idx_scale) const {
    return stride_batch_[idx_scale];
  }
};

template <int AccessSize, int MaxNumScales>
struct Params<half, Layout::CONCAT_NCHW32, AccessSize, MaxNumScales>
    : ParamsBase<half, 32, AccessSize, 128, MaxNumScales> {
  static constexpr Layout kSourceLayout = Layout::CONCAT_NCHW32;
  using Base = ParamsBase<half, 32, AccessSize, 128, MaxNumScales>;

  int stride_packed_;
  int stride_batch_;
  utils::Array<int, MaxNumScales> stride_height_;  // [S]

  utils::Array<half2, MaxNumScales> spatial_shapes_half2;  // [S,2]

  template <typename... Args>
  HOST_DEVICE_INLINE Params(Args... args) : Base(args...) {
    int value_offset = 0;
    for (int s = 0; s < this->NumScales; s++) {
      this->data_value[s] = reinterpret_cast<const void *>(
          reinterpret_cast<const half *>(this->data_value[0]) + value_offset);
      value_offset += this->spatial_shapes[s].x * this->spatial_shapes[s].y *
                      Base::kNumChannelsPacked;
      spatial_shapes_half2[s] =
          half2{this->spatial_shapes[s].x, this->spatial_shapes[s].y};
      stride_height_[s] = this->spatial_shapes[s].x * Base::kNumChannelsPacked;
    }
    stride_packed_ = value_offset;
    stride_batch_ = div_up(this->NumHeads * this->NumChannelsPerHead,
                           Base::kNumChannelsPacked) *
                    stride_packed_;
  }
  HOST_DEVICE_INLINE int stride_width(int idx_scale) const {
    return Base::kNumChannelsPacked;
  }
  HOST_DEVICE_INLINE int stride_height(int idx_scale) const {
    return stride_height_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_packed(int idx_scale) const {
    return stride_packed_;
  }
  HOST_DEVICE_INLINE int stride_batch(int idx_scale) const {
    return stride_batch_;
  }
};

template <int AccessSize, int MaxNumScales>
struct Params<half, Layout::CONCAT_NHWC, AccessSize, MaxNumScales>
    : ParamsBase<half, 64, AccessSize, 256, MaxNumScales> {
  static constexpr Layout kSourceLayout = Layout::CONCAT_NHWC;
  using Base = ParamsBase<half, 64, AccessSize, 256, MaxNumScales>;

  int stride_width_;
  int stride_batch_;
  utils::Array<int, MaxNumScales> stride_height_;  // [S]

  utils::Array<half2, MaxNumScales> spatial_shapes_half2;  // [S,2]

  template <typename... Args>
  HOST_DEVICE_INLINE Params(Args... args) : Base(args...) {
    stride_width_ = this->NumHeads * this->NumChannelsPerHead;
    int value_offset = 0;
    for (int s = 0; s < this->NumScales; s++) {
      this->data_value[s] = reinterpret_cast<const void *>(
          reinterpret_cast<const half *>(this->data_value[0]) + value_offset);
      value_offset +=
          this->spatial_shapes[s].x * this->spatial_shapes[s].y * stride_width_;
      spatial_shapes_half2[s] =
          half2{this->spatial_shapes[s].x, this->spatial_shapes[s].y};
      stride_height_[s] = this->spatial_shapes[s].x * stride_width_;
    }
    stride_batch_ = value_offset;
  }
  HOST_DEVICE_INLINE int stride_width(int idx_scale) const {
    return stride_width_;
  }
  HOST_DEVICE_INLINE int stride_height(int idx_scale) const {
    return stride_height_[idx_scale];
  }
  HOST_DEVICE_INLINE int stride_packed(int idx_scale) const {
    return Base::kNumChannelsPacked;
  }
  HOST_DEVICE_INLINE int stride_batch(int idx_scale) const {
    return stride_batch_;
  }
};
